#include <ntddk.h>

#include "KDevMon.h"
#include "AutoLock.h"
#include "DevMonManager.h"
#include "..\DevMonCommon\DevMonoCommon.h"

DRIVER_UNLOAD   DevMonUnload;
DRIVER_DISPATCH DevMonDeviceControl, HandleFilterFunction;

NTSTATUS
CompleteRequest(PIRP Irp, NTSTATUS Status = STATUS_SUCCESS, ULONG Info = 0);

NTSTATUS 
FilterAddDevice(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT PhysicalDeviceObject);

PCSTR 
MajorFunctionToString(UCHAR major);

DevMonManager g_Data;

EXTERN_C
NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegisterPath)
{
	UNICODE_STRING DevName      = RTL_CONSTANT_STRING(L"\\Device\\KDevMon");
	PDEVICE_OBJECT DeviceObject = nullptr;
	NTSTATUS Status             = STATUS_SUCCESS;

	Status = IoCreateDevice(DriverObject, 0, &DevName, FILE_DEVICE_UNKNOWN, 0, TRUE, &DeviceObject);
	if (NT_SUCCESS(Status) == FALSE) {
		KdPrint(("Device creation error (0x%08X)", Status));
		return Status;
	}

	UNICODE_STRING LinkName = RTL_CONSTANT_STRING(L"\\??\\KDevMon");
	Status = IoCreateSymbolicLink(&LinkName, &DevName);
	if (NT_SUCCESS(Status) == FALSE) {
		KdPrint(("Symbolic link creation error (0x%08X)", Status));
		IoDeleteDevice(DeviceObject);
		return Status;
	}

	DriverObject->DriverUnload = DevMonUnload;
	for (PDRIVER_DISPATCH& Func : DriverObject->MajorFunction) {
		Func = HandleFilterFunction;
	}

	g_Data.CDO = DeviceObject;
	g_Data.Init(DriverObject);

	return Status;
}

VOID
DevMonUnload(_In_ PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING LinkName = RTL_CONSTANT_STRING(L"\\??\\KDevMon");
	IoDeleteSymbolicLink(&LinkName);
	NT_ASSERT(g_Data.CDO);
	IoDeleteDevice(g_Data.CDO);

	g_Data.RemoveAllDevices();
}

NTSTATUS
CompleteRequest(PIRP Irp, NTSTATUS Status, ULONG Info)
{
	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = Info;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Status;
}

NTSTATUS 
FilterAddDevice(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT PhysicalDeviceObject)
{
	PDEVICE_OBJECT DeviceObject = nullptr;
	NTSTATUS Status = IoCreateDevice(DriverObject, 0, nullptr, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
	if (NT_SUCCESS(Status) == FALSE) {
		KdPrint(("Device creation error (0x%08X)", Status));
		return Status;
	}

	PDEVICE_OBJECT LowerDeviceObject = nullptr;
	Status = IoAttachDeviceToDeviceStackSafe(DeviceObject, PhysicalDeviceObject, &LowerDeviceObject);
	if (NT_SUCCESS(Status) == FALSE) {
		KdPrint(("Attached Device Stack Error (0x%08X)", Status));
		IoDeleteDevice(DeviceObject);
		return Status;
	}

	DeviceObject->Flags |= (LowerDeviceObject->Flags & (DO_BUFFERED_IO | DO_DIRECT_IO));
	DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;
	DeviceObject->Flags |= DO_POWER_PAGABLE;

	DeviceObject->DeviceType = LowerDeviceObject->DeviceType;

	return Status;
}

PCSTR
MajorFunctionToString(UCHAR major)
{
	static const char* Strings[] = {
		"IRP_MJ_CREATE",
		"IRP_MJ_CREATE_NAMED_PIPE",
		"IRP_MJ_CLOSE",
		"IRP_MJ_READ",
		"IRP_MJ_WRITE",
		"IRP_MJ_QUERY_INFORMATION",
		"IRP_MJ_SET_INFORMATION",
		"IRP_MJ_QUERY_EA",
		"IRP_MJ_SET_EA",
		"IRP_MJ_FLUSH_BUFFERS",
		"IRP_MJ_QUERY_VOLUME_INFORMATION",
		"IRP_MJ_SET_VOLUME_INFORMATION",
		"IRP_MJ_DIRECTORY_CONTROL",
		"IRP_MJ_FILE_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CONTROL",
		"IRP_MJ_INTERNAL_DEVICE_CONTROL",
		"IRP_MJ_SHUTDOWN",
		"IRP_MJ_LOCK_CONTROL",
		"IRP_MJ_CLEANUP",
		"IRP_MJ_CREATE_MAILSLOT",
		"IRP_MJ_QUERY_SECURITY",
		"IRP_MJ_SET_SECURITY",
		"IRP_MJ_POWER",
		"IRP_MJ_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CHANGE",
		"IRP_MJ_QUERY_QUOTA",
		"IRP_MJ_SET_QUOTA",
		"IRP_MJ_PNP",
	};
	NT_ASSERT(major <= IRP_MJ_MAXIMUM_FUNCTION);
	return Strings[major];
}

NTSTATUS 
HandleFilterFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	if (DeviceObject == g_Data.CDO) {
		switch (IoGetCurrentIrpStackLocation(Irp)->MajorFunction) {
		case IRP_MJ_CREATE:
		case IRP_MJ_CLOSE:
			return CompleteRequest(Irp);

		case IRP_MJ_DEVICE_CONTROL:
			return DevMonDeviceControl(DeviceObject, Irp);
		}
		return CompleteRequest(Irp, STATUS_INVALID_DEVICE_REQUEST);
	}

	DeviceExtension* Ext = (DeviceExtension*)DeviceObject->DeviceExtension;

	PETHREAD Thread = Irp->Tail.Overlay.Thread;
	HANDLE Tid      = nullptr;
	HANDLE Pid      = nullptr;
	if (Thread != nullptr) {
		Tid = PsGetThreadId(Thread);
		Pid = PsGetThreadProcessId(Thread);
	}

	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);

	KdPrint(("Driver: %wZ: Pid: %d, Tid: %ul, MJ = %d (%s)", 
		&Ext->m_LowerDeviceObject->DriverObject->DriverName, 
		HandleToULong(Pid), 
		HandleToULong(Tid), 
		Stack->MajorFunction,
		MajorFunctionToString(Stack->MajorFunction)));

	IoSkipCurrentIrpStackLocation(Irp);
	return IoCallDriver(Ext->m_LowerDeviceObject, Irp);
}

NTSTATUS 
DevMonDeviceControl(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS Status          = STATUS_INVALID_DEVICE_REQUEST;
	ULONG Code               = Stack->Parameters.DeviceIoControl.IoControlCode;

	switch (Code) {
	case IOCTL_DEVMON_ADD_DEVICE:
	case IOCTL_DEVMON_REMOVE_DEVICE: {
		WCHAR* Buffer = (WCHAR*)Irp->AssociatedIrp.SystemBuffer;
		ULONG Len = Stack->Parameters.DeviceIoControl.InputBufferLength;
		if (Buffer == nullptr || Len > 512) {
			Status = STATUS_INVALID_BUFFER_SIZE;
			break;
		}

		Buffer[Len / sizeof(WCHAR) - 1] = L'\0';
		if (Code == IOCTL_DEVMON_ADD_DEVICE) {
			Status = g_Data.AddDevice(Buffer);
		} else {
			Status = g_Data.RemoveDevice(Buffer) ? STATUS_SUCCESS : STATUS_NOT_FOUND;
		}
		break;
	}

	case IOCTL_DEVMON_REMOVE_ALL: {
		g_Data.RemoveAllDevices();
		Status = STATUS_SUCCESS;
		break;
	}
	}

	return CompleteRequest(Irp, Status);
}